from PyQt6.QtWidgets import (QApplication, QMessageBox, QInputDialog, QFileDialog,
                             QProgressDialog)
from PyQt6.QtGui import QPixmap, QDragEnterEvent, QDropEvent
from PyQt6.QtCore import Qt
import cv2
from paddleocr import PaddleOCR
import os
from datetime import datetime
from ui import BaseApp
from io import StringIO
import sys
from save_results import SaveResults

try:
    import pyi_splash
    pyi_splash.close()
except ImportError:
    pass

class MyApp(BaseApp):
    def __init__(self):
        super().__init__()
        self.ocr_engine = None
        self.current_image_path = None
        self.output_folder = os.path.expanduser("~/Desktop")  # 默认保存到桌面
        self.excel_file_name = "OCR_Results.xlsx"
        self.processing = False  # 防止重复点击
        # 新增属性
        self.image_paths = []  # 存储所有图片路径
        self.current_results = {}  # 存储识别结果 {路径: 结果文本}

        # 初始化OCR引擎
        self.init_ocr_engine()

        # 连接信号
        self.upload_btn.clicked.connect(self.open_image_dialog)
        self.ocr_btn.clicked.connect(self.run_ocr)
        self.select_folder_btn.clicked.connect(self.select_output_folder)
        self.input_filename_btn.clicked.connect(self.input_file_name)
        self.save_excel_btn.clicked.connect(self.save_results)
        self.lang_combo.currentIndexChanged.connect(self.init_ocr_engine)
        self.help_btn.clicked.connect(self.show_help)
        self.setAcceptDrops(True)
        self.setWindowTitle("OCR文字识别工具")
        # 连接新按钮信号
        self.add_folder_btn.clicked.connect(self.add_image_folder)
        self.clear_btn.clicked.connect(self.clear_image_list)
        self.image_list.itemClicked.connect(self.on_image_selected)
        # 初始化保存结果的类
        self.save_results_handler = SaveResults(self)

    def save_results(self):
        self.save_results_handler.save_results()

    def init_ocr_engine(self):
        lang = 'ch' if self.lang_combo.currentText() == "中文" else 'en'
        try:
            self.ocr_engine = PaddleOCR(use_angle_cls=True, lang=lang)
        except Exception as e:
            QMessageBox.critical(self, "引擎初始化失败", f"无法初始化OCR引擎：{str(e)}")

    def open_image_dialog(self):
        files, _ = QFileDialog.getOpenFileNames(
            self, "选择图片", "",
            "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if files:
            for path in files:
                if path not in self.image_paths:
                    self.image_paths.append(path)
                    self.image_list.addItem(path)

    # 清空列表
    def clear_image_list(self):
        self.image_paths.clear()
        self.image_list.clear()
        self.image_label.clear()
        self.current_image_path = None

    # 实现图片预览
    def on_image_selected(self, item):
        path = item.text()
        self.current_image_path = path
        self.display_image(path)

    def display_image(self, path):
        self.current_image_path = path
        pixmap = QPixmap(path)
        if not pixmap.isNull():
            scaled_pixmap = pixmap.scaled(
                self.image_label.width(),
                self.image_label.height(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.image_label.setPixmap(scaled_pixmap)
        else:
            self.show_error("图片加载失败", "不支持的文件格式或损坏的图片文件")

    def preprocess_image(self, image_path):
        try:
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError("无法读取图片文件")

            # 显示预处理进度
            progress = QProgressDialog("正在预处理图片...", None, 0, 0, self)
            progress.setWindowModality(Qt.WindowModality.WindowModal)
            progress.show()

            QApplication.processEvents()  # 更新UI

            # 预处理流程
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            binary = cv2.adaptiveThreshold(
                gray, 255,
                cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                cv2.THRESH_BINARY, 11, 2
            )
            temp_path = os.path.join(os.path.dirname(__file__), "temp_preprocessed.png")
            cv2.imwrite(temp_path, binary)

            progress.close()
            return temp_path

        except Exception as e:
            self.show_error("预处理失败", str(e))
            return None

    # 新增方法：添加文件夹
    def add_image_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
        if folder:
            self.find_images_in_folder(folder)

    # 新增方法：递归查找图片
    def find_images_in_folder(self, folder):
        extensions = ('.png', '.jpg', '.jpeg', '.bmp')
        for root, _, files in os.walk(folder):
            for file in files:
                if file.lower().endswith(extensions):
                    path = os.path.join(root, file)
                    if path not in self.image_paths:
                        self.image_paths.append(path)
                        self.image_list.addItem(path)

    # 修改OCR执行方法
    def run_ocr(self):
        if not self.image_paths:
            self.show_warning("提示", "请先添加要识别的图片")
            return

        self.result_area.clear()
        self.current_results.clear()

        progress = QProgressDialog("正在识别图片...", "取消", 0, len(self.image_paths), self)
        progress.setWindowModality(Qt.WindowModality.WindowModal)

        for i, path in enumerate(self.image_paths):
            progress.setValue(i)
            if progress.wasCanceled():
                break

            try:
                processed_path = self.preprocess_image(path)
                result = self.ocr_engine.ocr(processed_path, cls=True)

                # 处理空结果
                if not result or not result[0]:
                    error_msg = f"⚠️ 未识别到文字：{os.path.basename(path)}"
                    self.result_area.append(error_msg)
                    self.result_area.append("-" * 50 + "\n")
                    continue

                texts = [line[1][0] for line in result[0] if len(line) > 1]
                result_text = "\n".join(texts)

                self.current_results[path] = result_text
                self.result_area.append(f"✅ 识别成功：{os.path.basename(path)}")
                self.result_area.append(result_text)
                self.result_area.append("=" * 50 + "\n")

            except Exception as e:
                # 修改此处：将错误显示在结果区域而不是弹窗
                error_msg = f"""
    ❌ 识别失败：{os.path.basename(path)}
    错误类型：{type(e).__name__}
    错误详情：{str(e)}
                """
                self.result_area.append(error_msg)
                self.result_area.append("-" * 50 + "\n")

            finally:
                if 'processed_path' in locals() and os.path.exists(processed_path):
                    os.remove(processed_path)

        progress.close()
        # 恢复完成弹窗
        self.show_info("完成", "所有图片识别完成！")
        # 保留结果区域的提示
        self.result_area.append("\n▶ 识别任务已完成")

    # 修改后的拖放事件处理
    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QDropEvent):
        if self.processing:
            return

        # 获取所有拖放路径并过滤图片文件
        for url in event.mimeData().urls():
            path = url.toLocalFile()
            if os.path.isfile(path) and path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                if path not in self.image_paths:
                    self.image_paths.append(path)
                    self.image_list.addItem(path)
            elif os.path.isdir(path):  # 处理拖放文件夹
                self.find_images_in_folder(path)

        # 显示第一个图片预览
        if self.image_paths:
            self.display_image(self.image_paths[0])

    def select_output_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, "选择保存位置", self.output_folder)
        if folder_path:
            self.output_folder = folder_path
            QMessageBox.information(self, "路径已更新", f"文件将保存到：\n{folder_path}")

    def input_file_name(self):
        text, ok = QInputDialog.getText(
            self, "输入文件名",
            "请输入文件名（不包含扩展名）：",
            text=self.excel_file_name.split(".")[0]  # 默认文件名（去掉扩展名）
        )
        #文件名自动生成
        if text.strip() == "":  # 如果用户没有输入
            self.excel_file_name = f"OCR结果_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"

    # 辅助提示方法
    def show_error(self, title, message):
        QMessageBox.critical(self, title, message)

    def show_warning(self, title, message):
        QMessageBox.warning(self, title, message)

    def show_info(self, title, message):
        QMessageBox.information(self, title, message)

    def show_help(self):
        help_text = """
        <h3>OCR文字识别工具 - 使用指南</h3>
        <p><strong>上传图片</strong>: 点击“上传图片”按钮选择要识别的图片文件（支持拖动）。</p>
        <p><strong>选择语言</strong>: 在“识别语言”下拉菜单中选择图片文字的语言（中文或英文）。</p>
        <p><strong>开始识别</strong>: 点击“开始识别”按钮对上传的图片进行文字识别。</p>
        <p><strong>查看结果</strong>: 识别结果会显示在下方的文本框中。</p>
        <p><strong>选择文件夹</strong>: 点击“选择文件夹”按钮选择保存结果的文件夹。</p>
        <p><strong>输入文件名</strong>: 点击“输入文件名”按钮为保存的文件命名。</p>
        <p><strong>保存结果</strong>: 点击“保存结果”按钮将识别结果保存为Excel或文本文件。</p>
        """
        QMessageBox.about(self, "帮助", help_text)

    if sys.stderr is None:
        sys.stderr = StringIO()

if __name__ == "__main__":
    import sys
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 使用更现代的主题
    window = MyApp()
    window.show()
    sys.exit(app.exec())
